package learner.hmm;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import learner.ModelLearner;
import learner.Pattern;
import learner.PatternModel;

public class BaumWelchLearner implements ModelLearner {

	public static int NB_OF_STATES = 3;
	
	private Matrix<Integer, Integer> a ;
	private Matrix<Integer, Integer> newA;
	
	private Matrix<Integer, MovementVector> b;
	private Matrix<Integer, MovementVector> newB;
	
	private Map<Integer, Double> pi;
	
	private AlgorithmValueCalculator valueCalculator;

	private int counter;
	
	@Override
	public PatternModel learnModel(String name, List<Pattern> patterns) {
		MovementStats movementStatistics = new MovementStats(patterns);
		a = new Matrix<Integer, Integer>(initializeA());
		b = new Matrix<Integer, MovementVector>(initializeB(movementStatistics));
		if (!hasValidB())
			System.err.print("b invalid");
		pi = initializePi();
		valueCalculator = new  AlgorithmValueCalculator(a, b, pi, movementStatistics.getMovement(2)); 
		// TODO: convergentie uitbreiden zodat het alle paden van het patroon gebruikt.
//		do {
//			valueCalculator = new  AlgorithmValueCalculator(a, b, pi, movementStatistics.getMovement(0)); 
//			improveA(movementStatistics.getMovement(0));
//			improveB(movementStatistics.getMovement(0));
//			improvePi(movementStatistics.getMovement(0));
//			a = newA;
//			if (!hasValidA())
//				System.out.println("A invalid");;
//			b = newB;
//		} while (! localMaxReached(newA, a, newB, b));
		System.out.println(valueCalculator.getChanceOnPattern());
		return new HMMModel(name, a.toMap(), b.toMap(), pi);
	}
	
	private boolean hasValidB() {
		for (int i = 0; i < NB_OF_STATES ; i ++) {
			if (b.sumRow(i) > (double)1 + 0.1)
				return false;
		}
		return true;
	}

	public boolean hasValidA() {
		for (int i = 0; i < NB_OF_STATES ; i ++) {
			if (a.sumRow(i) > (double)1 + 0.00000001)
				return false;
		}
		return true;
	}

	

	private Map<Integer, Map<Integer, Double>> initializeA() {
		Map<Integer, Map<Integer, Double>> a = new HashMap<Integer, Map<Integer, Double>>();
		for (int i = 0 ; i < NB_OF_STATES ; i++) {
			Map<Integer, Double> probabilities = calculateInitialTranferProbabilities(i);
			a.put(i, probabilities);
		}
		return a;
	}
	
	private Map<Integer, Double> calculateInitialTranferProbabilities(int state) {
		if ( state >= NB_OF_STATES)
			throw new IllegalArgumentException();
		Map<Integer, Double> probabilities =  new HashMap<Integer, Double>();
		if (state == NB_OF_STATES - 1)
			probabilities.put(state, 1.0);
		else {
			probabilities.put(state, 0.9);
			probabilities.put(state + 1, 0.1);
		}
		return probabilities;
	}
	
	private Map<Integer, Map<MovementVector, Double>> initializeB(MovementStats movementStatistics) {
		List<Map<MovementVector, Double>> conditionalProbs = movementStatistics.getRelativeFreqMovementIntervals(NB_OF_STATES);
		Map<Integer, Map<MovementVector, Double>> contionalProbsMap = new HashMap<Integer, Map<MovementVector,Double>>(); 
		for (int i = 0 ;  i < conditionalProbs.size(); i ++) {
			contionalProbsMap.put(i , conditionalProbs.get(i));
		}
		return contionalProbsMap;
	}

	private Map<Integer, Double> getConditionalProbabilitiesFor(
			Map<MovementVector, Double> relativeFrequenciesMovements,
			MovementVector movement) {
		Map<Integer, Double> conditionalProbabilities = new HashMap<Integer, Double>();
		for (int i = 0 ; i < NB_OF_STATES ; i ++) {
			conditionalProbabilities.put(i, relativeFrequenciesMovements.get(movement));
		}
		return conditionalProbabilities;
	}

	
	private Map<Integer, Double> initializePi() {
		Map<Integer, Double> pi = new HashMap<Integer, Double>();
		pi.put(0,1.0);
		for (int i = 1 ; i < NB_OF_STATES ; i ++) {
			pi.put(i, 0.0);
		}
		return pi;
	}

	private void improveA(Movement movement) {
		newA = new Matrix<Integer, Integer>(new HashMap<Integer, Map<Integer, Double>>());
		for (int i = 0 ; i < NB_OF_STATES ; i++){
			for (int j = 0 ; j < NB_OF_STATES; j++) {
				double numerator = 0 ;
				double denominator = 0;
				for (int k = 1; k < movement.getSize() ;  k++ ) {
					numerator += valueCalculator.getPsiForTime(k, i, j);
					denominator += valueCalculator.getGammaForTime(k, i);		
				}
				newA.setElement(i, j , numerator / denominator); 
			}	
		}
		a = newA;
	}
	
	private void improveB(Movement movement) {
		newB = new Matrix<Integer, MovementVector>();
		for (int j = 0 ; j < NB_OF_STATES ; j ++) {
			for (MovementVector vm : movement.getDifferentMovements()) {
				double numerator = getNbOfExpectedMovIn(j, vm, movement);
				double denominator = getExpecteTimesIn(j, movement);
				newB.setElement(j, vm, numerator / denominator);
			}
		}
	}

	
	private double getExpecteTimesIn(int j, Movement movement) {
		double expectedTimesInJ = 0.0;
		for (int k = 1 ; k <= movement.getSize() ; k ++) {
			expectedTimesInJ += valueCalculator.getGammaForTime(k, j);
		}
		return expectedTimesInJ;
	}

	private double getNbOfExpectedMovIn(int j, MovementVector vm,
			Movement movement) {
		double expectedMov = 0;
		for (int k = 1 ; k <= movement.getSize() ; k ++) {
			if (movement.getMovementForTime(k).equals(vm))
				expectedMov += valueCalculator.getGammaForTime(k,j);
		}
		return expectedMov;
	}

	private void improvePi(Movement movement) {
		for (int i = 0 ; i < NB_OF_STATES ; i ++) {
			pi.put(i,valueCalculator.getGammaForTime(1, i));
		}
	}
	
	private boolean localMaxReached(Matrix<Integer, Integer> oldA,
			Matrix<Integer, Integer> a,
			Matrix<MovementVector, Integer> oldB,
			Matrix<MovementVector, Integer> b) {
		counter ++;
		return counter == 20;
	}


}
